MODULE completeness_profile
  !-----------------------------------------------------------
  ! Examines the completeness of one-electron basis set via 
  ! overlap integrals <basis | scan (alpha)>, where a basis  
  ! consists of contracted GTO's and scan (alpha) are the trial 
  ! GTOs with varying exponent alpha. 
  !-----------------------------------------------------------
  ! See CHANGELOG
  !-----------------------------------------------------------
  ! Copyright (C) 2003-2008 Pekka Manninen, 2010 Jussi Lehtola
  !
  ! This program is distributed under the terms of the GNU General
  ! Public License as published by the Free Software Foundation;
  ! either version 2 of the License, or (at your option) any later version.
  !
  ! This program is distributed in the hope that it will be useful,
  ! but WITHOUT ANY WARRANTY; without even the implied warranty of
  ! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. 
  !
  ! Please note that this program as well as its 
  ! author must be properly cited whenever the program or some parts 
  ! originated on it are employed. 
  !
  !-----------------------------------------------------------
  USE definitions
  USE error
  IMPLICIT NONE
  PRIVATE
  PUBLIC :: comp_prof, get_alphas, normalize_contraction
  !-----------------------------------------------------------
CONTAINS
  !-----------------------------------------------------------

  SUBROUTINE get_alphas(mina, maxa, alphas)
    ! Compute table of exponents

    ! Minimum and maximum value of exponent
    REAL(KIND=prec), INTENT(IN) :: mina, maxa
    REAL(KIND=prec), DIMENSION(:), INTENT(OUT) :: alphas

    ! Logarithms
    REAL(KIND=prec) :: lgmin, lgmax
    ! Step size
    REAL(KIND=prec) :: step

    ! Number of points
    INTEGER :: np
    ! Loop index
    INTEGER :: ia

    ! Get values
    lgmin=LOG10(mina)
    lgmax=LOG10(maxa)
    np=SIZE(alphas)
    
    ! Logarithmic step size for scanning exponent
    step = (lgmax - lgmin)/REAL((np-1),prec)
    ! Fill G(alpha) exponent array
    DO ia=1, np
       alphas(ia)=10.0_prec**(lgmin + (ia-1.0_prec)*step)
    END DO

  END SUBROUTINE get_alphas


  SUBROUTINE comp_prof(l, X, C, alphas, usesto, Y)
    IMPLICIT NONE
    INTRINSIC :: SQRT

    ! Value of l to compute profile for
    INTEGER, INTENT(IN) :: l
    ! Exponents of basis set
    REAL(KIND=prec), DIMENSION(:), INTENT(IN) :: X
    ! Contraction coefficients, these should be normalized
    REAL(KIND=prec), DIMENSION(:,:), INTENT(IN) :: C

    ! Values of alpha to scan with
    REAL(KIND=prec), DIMENSION(:), INTENT(IN) :: alphas
    ! Scan with an STO?
    LOGICAL, INTENT(IN) :: usesto

    ! Completeness profile
    REAL(KIND=prec), DIMENSION(SIZE(alphas)), INTENT(OUT) :: Y

    ! Overlaps with scanning function
    REAL(KIND=prec), DIMENSION(SIZE(X),SIZE(alphas)) :: ovl
    ! Overlap of scanning function with basis set
    REAL(KIND=prec), DIMENSION(SIZE(C,2),SIZE(alphas)) :: J
    ! Helper array
    REAL(KIND=prec), DIMENSION(SIZE(X),SIZE(C,2)) :: K
    ! Helper arrays for S and inverse of Cholesky decomposition
    REAL(KIND=prec), DIMENSION(SIZE(C,2),SIZE(C,2)) :: S, T
    ! Boolean for success of inversion of S
    LOGICAL :: invsucc

    ! Numbers of exponents, contractions and alpha points
    INTEGER :: nx, nc, na

    ! Helper for completeness profile
    REAL(KIND=prec) :: jma
    ! Loop indices
    INTEGER :: ia, m

    nx=SIZE(X)
    nc=SIZE(C,2)
    na=SIZE(alphas)

    IF(nx /= SIZE(C,1)) THEN
       CALL error_stop('Error in completeness_profile.f90 - exponents and contraction coefficients do not match!')
    END IF

    ! Form overlap matrix of contracted basis set
    CALL overlap(l,nx,nc,X,C,S)

    ! Perform Cholesky orthonormalization of contracted basis set to form S^(-1/2) in T
    CALL invcholeskym(nc,S,T,invsucc)

    ! If inversion failed, set profile to zero and return.
    IF ( .NOT. invsucc ) THEN
       WRITE (*,*) 'Unable to orthogonalize, reduce number of exponents if necessary'
       Y = 0.0_prec
       RETURN
    END IF

    ! Get scan overlaps
    CALL scanoverlap(l,X,alphas,usesto,ovl)

    ! Contract coefficients with T^(-1), where S=T T^T
!    K = MATMUL(C(1:nx,1:nc),T) !  (nx,nc) = (nx,nc) x (nc,nc)
    CALL matm('n',C,'n',T,K)

    ! Contract the result with the overlap with the scanning functions
!    J = MATMUL(TRANSPOSE(K),invl) ! (nc,points) = (nc,nx) x (nx,points)
    CALL matm('t',K,'n',ovl,J)

    ! Calculate completeness profile
    Y = 0.0_prec
    DO ia = 1, SIZE(alphas)
       DO m = 1, nc
          jma = J(m,ia)
          Y(ia) = Y(ia) + jma * jma
       END DO
    END DO
  END SUBROUTINE comp_prof


  SUBROUTINE scanoverlap(l, X, alphas, usesto, ovl)
    ! Calculate overlaps of scanning exponents and primitives

    ! C function that calculates overlap of normalized STO and normalized GTO
    REAL(KIND=prec) :: sto_gto_overlap

    ! Value of l to compute profile for
    INTEGER, INTENT(IN) :: l
    ! Exponents of basis set
    REAL(KIND=prec), DIMENSION(:), INTENT(IN) :: X
    ! Values of alpha to scan with
    REAL(KIND=prec), DIMENSION(:), INTENT(IN) :: alphas
    ! Scan with an STO?
    LOGICAL, INTENT(IN) :: usesto

    ! Loop indices
    INTEGER :: ia, ix
    ! and upper limits
    INTEGER :: na, nx

    ! Overlaps
    REAL(KIND=prec), DIMENSION(SIZE(X),SIZE(alphas)) :: ovl

    IF(.NOT. usesto) THEN

       ! Get overlaps of scanning exponents and primitives
       CALL primitive_overlap(l,X,alphas,ovl)

    ELSE

       na=SIZE(alphas)
       nx=SIZE(X)

       ! Use an STO to scan
       DO ia = 1, na ! Loop over scan exponents
          ! Loop over exponents in basis
          DO ix = 1, nx
             ! Scan GTO basis set with an STO, take Fortan indexing of l into account
             ovl(ix,ia)=sto_gto_overlap(alphas(ia),X(ix),l-1)
          END DO
       END DO
    END IF
  END SUBROUTINE scanoverlap

  !------------------------------------------------------------
  ! Needed routines
  !------------------------------------------------------------
  SUBROUTINE overlap(l,nx,nc,chi,cont,ovl)
    ! overlap matrix for a given block of l-type functions
    IMPLICIT NONE
    INTRINSIC SQRT
    !  REAL(KIND=prec) gto_gto_overlapint ! function in overlap.c
    EXTERNAL initoverlap ! function in overlap.c
    INTEGER, INTENT(IN) :: l, nx, nc
    REAL(KIND=prec), INTENT(IN), DIMENSION(nx) :: chi
    REAL(KIND=prec), INTENT(IN), DIMENSION(nx,nc) :: cont
    REAL(KIND=prec), INTENT(OUT), DIMENSION(nc,nc) :: ovl

    ! Overlaps of unnormalized primitives
    REAL(KIND=prec), DIMENSION(nx,nx) :: povl
    ! Helper array
    REAL(KIND=prec), DIMENSION(nc,nx) :: A

    ! Check here that overlap table is initialized
    call initoverlap()

    ! Get overlaps of primitives
    call primitive_overlap(l,chi,chi,povl)

    ! Contract right-hand side with contraction coefficients  
!    A = MATMUL(TRANSPOSE(cont),povl)
    CALL matm('t',cont,'n',povl,A)

    ! left-hand side, too, to end up with final result
!    ovl = MATMUL(A,cont)
    CALL matm('n',A,'n',cont,ovl)

  END SUBROUTINE overlap
  !--------------------------------------------------------------
  SUBROUTINE primitive_overlap(l,chia,chib,ovl)
    ! Compute overlaps of normalized Gaussian primitives

    ! Angular momentum and number of exponents
    INTEGER, INTENT(IN) :: l
    ! First set of exponents
    REAL(KIND=prec), INTENT(IN), DIMENSION(:) :: chia
    ! Second set of exponents
    REAL(KIND=prec), INTENT(IN), DIMENSION(:) :: chib
    ! Overlaps
    REAL(KIND=prec), INTENT(OUT), DIMENSION(SIZE(chia),SIZE(chib)) :: ovl

    ! Helpers
    REAL(KIND=prec) :: chiij, tmp
    ! Square root and fourth root of tmp
    REAL(KIND=prec) :: tmp2, tmp4

    ! Sizes of arrays
    INTEGER :: na, nb
    ! Loop indices
    INTEGER :: ia, ib

    na=SIZE(chia)
    nb=SIZE(chib)

    IF (l==1) THEN ! S type
       DO ia = 1, na
          DO ib = 1, nb
             ! Calculate sum
             chiij=chia(ia)+chib(ib)
             ! Calculate helper
             tmp=4*chia(ia)*chib(ib)/(chiij*chiij)
             ! Need tmp^(3/4)
             tmp2=SQRT(tmp)
             tmp4=SQRT(tmp2)
             ovl(ia,ib)=tmp2*tmp4
          END DO  
       END DO
    ELSE IF (l==2) THEN ! P type
       DO ia = 1, na
          DO ib = 1, nb
             ! Calculate sum
             chiij=chia(ia)+chib(ib)
             ! Calculate helper
             tmp=4*chia(ia)*chib(ib)/(chiij*chiij)
             ! Need tmp^(5/4)
             tmp2=SQRT(tmp)
             tmp4=SQRT(tmp2)
             ovl(ia,ib)=tmp*tmp4
          END DO
       END DO       
    ELSE
       DO ia = 1, na
          DO ib = 1, nb
             ! Calculate sum
             chiij=chia(ia)+chib(ib)
             ! Calculate helper
             tmp=4*chia(ia)*chib(ib)/(chiij*chiij)
             ! Need tmp^(l/2+3/4). Remember that the real l is Fortran l-1.
             ovl(ia,ib)=tmp**((l-1)/2.0_prec+0.75_prec)
          END DO
       END DO
    END IF
  END SUBROUTINE primitive_overlap
  !--------------------------------------------------------------
  SUBROUTINE normalize_contraction(l,chi,cont,normcont)
    ! Normalize the contractions

    ! Angular momentum
    INTEGER, INTENT(IN) :: l
    ! Exponents
    REAL(KIND=prec), INTENT(IN), DIMENSION(:) :: chi
    ! Contraction coefficients
    REAL(KIND=prec), INTENT(IN), DIMENSION(:,:) :: cont
    ! Normalized contractions
    REAL(KIND=prec), INTENT(OUT), DIMENSION(:,:) :: normcont    

    ! Overlaps of primitives
    REAL(KIND=prec), DIMENSION(SIZE(chi),SIZE(chi)) :: povl

    ! Get overlaps of primitives
    CALL primitive_overlap(l,chi,chi,povl)
    
    ! Now we can compute the normalization
    CALL normalize_contraction_ovl(povl,cont,normcont)
  END SUBROUTINE normalize_contraction

  ! Worker routine
  SUBROUTINE normalize_contraction_ovl(povl,cont,normcont)
    ! Normalize the contraction, given the overlaps of primitives

    ! Overlaps of primitives
    REAL(KIND=prec), DIMENSION(:,:) :: povl
    ! Contraction coefficients
    REAL(KIND=prec), INTENT(IN), DIMENSION(:,:) :: cont
    ! Normalized contractions
    REAL(KIND=prec), INTENT(OUT), DIMENSION(SIZE(cont,1),SIZE(cont,2)) :: normcont    

    ! Number of primitives and contractions
    INTEGER :: nx, nc

    ! Loop indices
    INTEGER :: j

    ! Helper
    INTEGER :: ic
    REAL(KIND=prec) :: norm
    REAL(KIND=prec), DIMENSION(SIZE(povl,1)) :: cnorm

    nx=SIZE(cont,1)
    nc=SIZE(cont,2)

    WRITE(*,*) 'Normalize_contraction_ovl called.'

    ! FIXME TODO - Use matrix multiplication here
    DO ic=1,nc
       ! Reset norm
       norm=0.0_prec
!       ! Normalization coefficient is
       DO j=1,nx
!          DO k=1,nx
!             norm=norm+cont(j,ic)*povl(j,k)*cont(k,ic)
!          END DO
          cnorm=matmul(povl,cont(:,ic))
          norm=dot_product(cont(:,ic),cnorm)

       END DO
       
       ! and we normalize with
       normcont(ic,:)=cont(ic,:)/SQRT(norm)
    END DO
  END SUBROUTINE normalize_contraction_ovl
  !--------------------------------------------------------------
  SUBROUTINE invcholeskym(n,S,U,succ)
    ! Inverse Cholesky matrix of a given matrix S
    ! (U^-1 of S = U^T U) 
    IMPLICIT NONE
    EXTERNAL SPOTRF, DPOTRF, STRTRI, DTRTRI
    INTEGER, INTENT(IN) :: n
    INTEGER :: info
    REAL(KIND=prec), INTENT(IN), DIMENSION(n,n) :: S
    REAL(KIND=prec), INTENT(OUT), DIMENSION(n,n) :: U
    REAL(KIND=prec), DIMENSION(n,n) :: T
    ! Was the invertion a success?
    LOGICAL, INTENT(OUT) :: succ

    ! Default value:
    succ = .TRUE.

    ! Compute S = T T^T
    CALL cholesky(n,S,T,info)
    IF (info == 0) THEN
       ! Form T^-1 in U
       CALL tri_invert(n,T,U,info) 
       IF (info == 0) THEN
          RETURN
       ELSE
          ! CALL error_stop('Unable to construct the U^-1 matrix')
          succ = .FALSE.
       END IF
    ELSE
       ! CALL error_stop('Unable to construct the U matrix')
       succ = .FALSE.
    END IF
  END SUBROUTINE invcholeskym
  !--------------------------------------------------------------
  ! LAPACK interfaces
  !--------------------------------------------------------------
  SUBROUTINE matm(ta,A,tb,B,C)
    ! Calculate matrix product AB using LAPACK

    IMPLICIT NONE
    ! LAPACK functions
    EXTERNAL SGEMM, DGEMM

    ! Transpose matrix A before multiplication?
    CHARACTER(LEN=1), INTENT(IN) :: ta
    ! Transpose matrix B before multiplication?
    CHARACTER(LEN=1), INTENT(IN) :: tb
    ! The input matrices
    REAL(KIND=prec), DIMENSION(:,:), INTENT(IN) :: A, B
    ! The result of the multiplication
    REAL(KIND=prec), DIMENSION(:,:), INTENT(OUT) :: C

    ! A is ma x na, B is mb x nb, result is mc x nc
    INTEGER :: ma, na, mb, nb, mc, nc

    ! Number of elements in product sum
    INTEGER :: pa=0, pb=0

    ! Get sizes
    ma=SIZE(A,1)
    na=SIZE(A,2)

    mb=SIZE(B,1)
    nb=SIZE(B,2)

    mc=SIZE(C,1)
    nc=SIZE(C,2)


    ! Check that arrays are correctly shaped for product
    IF(ta=='n' .OR. ta=='N') THEN
       ! Don't transpose A - sum is over na
       pa=na
       IF(mc/=ma) THEN
          CALL error_stop('mc /= ma in completeness_profile.f90')
       END IF
    ELSE IF(ta=='t' .OR. ta=='T') THEN
       ! Transpose - sum is over ma
       pa=ma
       IF(mc/=na) THEN
          CALL error_stop('mc /= na in completeness_profile.f90')
       END IF
    END IF

    IF(tb=='n' .OR. tb=='N') THEN
       ! Don't transpose B - sum is over mb
       pb=mb
       IF(nc/=nb) THEN
          CALL error_stop('nc /= nb in completeness_profile.f90')
       END IF
    ELSE IF(tb=='t' .OR. tb=='T') THEN
       pb=nb
       IF(nc/=mb) THEN
          CALL error_stop('nc /= mb in completeness_profile.f90')
       END IF
    END IF

    IF(pa /= pb) THEN
       CALL error_stop('Arrays are inconsistently shaped for matrix product.')
    END IF

    ! Result: C= alpha*AB + beta*C
    C=0.0_prec
    IF (prec < 8) THEN
       ! alpha = 1.0, beta=0.0
       CALL SGEMM(ta,tb,mc,nc,pa,1.0_prec,A,ma,B,mb,0.0_prec,C,mc)
    ELSE
       CALL DGEMM(ta,tb,mc,nc,pa,1.0_prec,A,ma,B,mb,0.0_prec,C,mc)
    END IF
  END SUBROUTINE MATM
  !---------------------------------------------------------------------
  SUBROUTINE cholesky(n,a,b,info)
    ! Cholesky decomposition a = b^T b
    IMPLICIT NONE
    EXTERNAL SPOTRF, DPOTRF
    INTEGER, INTENT(IN) :: n
    INTEGER, INTENT(OUT) :: info
    INTEGER :: i
    REAL(KIND=prec), DIMENSION(n,n), INTENT(IN) :: a
    REAL(KIND=prec), DIMENSION(n,n), INTENT(OUT) :: b
    b = a
    IF (prec < 8) THEN
       CALL SPOTRF('U', n, b, n, info)
    ELSE
       CALL DPOTRF('U', n, b, n, info)
    END IF
    IF (info == 0) THEN
       DO i = 1, n
          b(i+1:n,i) = 0.0_prec
       END DO
       RETURN
    ELSE
       !     WRITE (*,'(A,I3)') &
       !          'Warning: xPOTRF (cholesky in completeness_profile.f90) failed, ifail=',&
       !          info
       !     DO i = 1, n
       !        WRITE(*,'(8e10.2)') a(i,:)
       !     END DO
    END IF
  END SUBROUTINE cholesky
  !---------------------------------------------------------------------
  SUBROUTINE tri_invert(n,a,b,info) 
    ! inverse of a triangular matrix
    IMPLICIT NONE
    EXTERNAL STRTRI, DTRTRI
    INTEGER, INTENT(IN) :: n
    INTEGER, INTENT(OUT) :: info
    REAL(KIND=prec), DIMENSION(n,n), INTENT(IN) :: a
    REAL(KIND=prec), DIMENSION(n,n), INTENT(OUT) :: b
    b = a
    IF (prec < 8) THEN
       CALL STRTRI('U', 'N', n, b, n, info)
    ELSE
       CALL DTRTRI('U', 'N', n, b, n, info)
    END IF
    IF (info == 0) THEN
       RETURN
    ELSE
       WRITE (*,'(A,I3)') &
            'Warning: xTRTRI (tri_invert in completeness_profile.f90) failed, ifail=',&
            info
    END IF
  END SUBROUTINE tri_invert

END MODULE completeness_profile
